{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute this cell to install dependencies\n",
    "%pip install sf-hamilton[visualization]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# First-class Caching in Hamilton [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dagworks-inc/hamilton/blob/main/examples/caching_nodes/caching.ipynb) [![GitHub badge](https://img.shields.io/badge/github-view_source-2b3137?logo=github)](https://github.com/apache/hamilton/blob/main/examples/caching_nodes/caching.ipynb)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "from hamilton import registry\n",
    "registry.disable_autoload()\n",
    "registry.load_extension(\"pandas\")\n",
    "\n",
    "%load_ext hamilton.plugins.jupyter_magic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 2.43.0 (0)\n",
       " -->\n",
       "<!-- Title: %3 Pages: 1 -->\n",
       "<svg width=\"469pt\" height=\"304pt\"\n",
       " viewBox=\"0.00 0.00 468.50 304.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 300)\">\n",
       "<title>%3</title>\n",
       "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-300 464.5,-300 464.5,4 -4,4\"/>\n",
       "<g id=\"clust1\" class=\"cluster\">\n",
       "<title>cluster__legend</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"black\" points=\"8,-156 8,-288 104,-288 104,-156 8,-156\"/>\n",
       "<text text-anchor=\"middle\" x=\"56\" y=\"-272.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">Legend</text>\n",
       "</g>\n",
       "<!-- G -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>G</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M93.5,-64C93.5,-64 18.5,-64 18.5,-64 12.5,-64 6.5,-58 6.5,-52 6.5,-52 6.5,-12 6.5,-12 6.5,-6 12.5,0 18.5,0 18.5,0 93.5,0 93.5,0 99.5,0 105.5,-6 105.5,-12 105.5,-12 105.5,-52 105.5,-52 105.5,-58 99.5,-64 93.5,-64\"/>\n",
       "<text text-anchor=\"start\" x=\"50\" y=\"-42.8\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">G</text>\n",
       "<text text-anchor=\"start\" x=\"17.5\" y=\"-14.8\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">DataFrame</text>\n",
       "</g>\n",
       "<!-- B -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>B</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M188,-146C188,-146 158,-146 158,-146 152,-146 146,-140 146,-134 146,-134 146,-94 146,-94 146,-88 152,-82 158,-82 158,-82 188,-82 188,-82 194,-82 200,-88 200,-94 200,-94 200,-134 200,-134 200,-140 194,-146 188,-146\"/>\n",
       "<text text-anchor=\"start\" x=\"167.5\" y=\"-124.8\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">B</text>\n",
       "<text text-anchor=\"start\" x=\"157.5\" y=\"-96.8\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">float</text>\n",
       "</g>\n",
       "<!-- E -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>E</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M365.5,-176C365.5,-176 335.5,-176 335.5,-176 329.5,-176 323.5,-170 323.5,-164 323.5,-164 323.5,-124 323.5,-124 323.5,-118 329.5,-112 335.5,-112 335.5,-112 365.5,-112 365.5,-112 371.5,-112 377.5,-118 377.5,-124 377.5,-124 377.5,-164 377.5,-164 377.5,-170 371.5,-176 365.5,-176\"/>\n",
       "<text text-anchor=\"start\" x=\"345.5\" y=\"-154.8\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">E</text>\n",
       "<text text-anchor=\"start\" x=\"341\" y=\"-126.8\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">str</text>\n",
       "</g>\n",
       "<!-- B&#45;&gt;E -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>B&#45;&gt;E</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M200.38,-117.84C224.83,-121.48 262.19,-127.22 294.5,-133 300.52,-134.08 306.89,-135.29 313.09,-136.5\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"312.75,-140 323.24,-138.52 314.12,-133.14 312.75,-140\"/>\n",
       "</g>\n",
       "<!-- C -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>C</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M282.5,-206C282.5,-206 252.5,-206 252.5,-206 246.5,-206 240.5,-200 240.5,-194 240.5,-194 240.5,-154 240.5,-154 240.5,-148 246.5,-142 252.5,-142 252.5,-142 282.5,-142 282.5,-142 288.5,-142 294.5,-148 294.5,-154 294.5,-154 294.5,-194 294.5,-194 294.5,-200 288.5,-206 282.5,-206\"/>\n",
       "<text text-anchor=\"start\" x=\"262\" y=\"-184.8\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">C</text>\n",
       "<text text-anchor=\"start\" x=\"252.5\" y=\"-156.8\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">bool</text>\n",
       "</g>\n",
       "<!-- B&#45;&gt;C -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>B&#45;&gt;C</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M200.26,-131.03C210.11,-137.42 221.45,-144.78 231.94,-151.58\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"230.2,-154.63 240.5,-157.13 234.01,-148.75 230.2,-154.63\"/>\n",
       "</g>\n",
       "<!-- F -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>F</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M448.5,-176C448.5,-176 418.5,-176 418.5,-176 412.5,-176 406.5,-170 406.5,-164 406.5,-164 406.5,-124 406.5,-124 406.5,-118 412.5,-112 418.5,-112 418.5,-112 448.5,-112 448.5,-112 454.5,-112 460.5,-118 460.5,-124 460.5,-124 460.5,-164 460.5,-164 460.5,-170 454.5,-176 448.5,-176\"/>\n",
       "<text text-anchor=\"start\" x=\"428.5\" y=\"-154.8\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">F</text>\n",
       "<text text-anchor=\"start\" x=\"420.5\" y=\"-126.8\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">dict</text>\n",
       "</g>\n",
       "<!-- E&#45;&gt;F -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>E&#45;&gt;F</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M377.57,-144C383.5,-144 389.9,-144 396.16,-144\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"396.47,-147.5 406.47,-144 396.47,-140.5 396.47,-147.5\"/>\n",
       "</g>\n",
       "<!-- A -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>A</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M71,-146C71,-146 41,-146 41,-146 35,-146 29,-140 29,-134 29,-134 29,-94 29,-94 29,-88 35,-82 41,-82 41,-82 71,-82 71,-82 77,-82 83,-88 83,-94 83,-94 83,-134 83,-134 83,-140 77,-146 71,-146\"/>\n",
       "<text text-anchor=\"start\" x=\"50.5\" y=\"-124.8\" font-family=\"Helvetica,sans-Serif\" font-weight=\"bold\" font-size=\"14.00\">A</text>\n",
       "<text text-anchor=\"start\" x=\"46.5\" y=\"-96.8\" font-family=\"Helvetica,sans-Serif\" font-style=\"italic\" font-size=\"14.00\">int</text>\n",
       "</g>\n",
       "<!-- A&#45;&gt;B -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>A&#45;&gt;B</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M83.1,-114C98.57,-114 118.43,-114 135.44,-114\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"135.61,-117.5 145.61,-114 135.61,-110.5 135.61,-117.5\"/>\n",
       "</g>\n",
       "<!-- C&#45;&gt;E -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>C&#45;&gt;E</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M294.57,-164.34C300.79,-162.04 307.53,-159.55 314.08,-157.12\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"315.31,-160.39 323.47,-153.64 312.88,-153.83 315.31,-160.39\"/>\n",
       "</g>\n",
       "<!-- _C_inputs -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>_C_inputs</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"211.5,-209.5 134.5,-209.5 134.5,-164.5 211.5,-164.5 211.5,-209.5\"/>\n",
       "<text text-anchor=\"start\" x=\"150\" y=\"-182.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">D</text>\n",
       "<text text-anchor=\"start\" x=\"167\" y=\"-182.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">bool</text>\n",
       "</g>\n",
       "<!-- _C_inputs&#45;&gt;C -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>_C_inputs&#45;&gt;C</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M211.62,-181.71C217.78,-180.85 224.14,-179.95 230.26,-179.09\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"230.83,-182.55 240.24,-177.69 229.85,-175.62 230.83,-182.55\"/>\n",
       "</g>\n",
       "<!-- input -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>input</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"85.5,-256.5 26.5,-256.5 26.5,-219.5 85.5,-219.5 85.5,-256.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"56\" y=\"-234.3\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input</text>\n",
       "</g>\n",
       "<!-- function -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>function</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M84,-201.5C84,-201.5 28,-201.5 28,-201.5 22,-201.5 16,-195.5 16,-189.5 16,-189.5 16,-176.5 16,-176.5 16,-170.5 22,-164.5 28,-164.5 28,-164.5 84,-164.5 84,-164.5 90,-164.5 96,-170.5 96,-176.5 96,-176.5 96,-189.5 96,-189.5 96,-195.5 90,-201.5 84,-201.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"56\" y=\"-179.3\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">function</text>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7fa3d5bc4150>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%cell_to_module primitives -d\n",
    "import pandas as pd\n",
    "from hamilton.function_modifiers import tag\n",
    "\n",
    "def A() -> int:\n",
    "    return 7\n",
    "\n",
    "def B(A: int) -> float:\n",
    "    return float(A)\n",
    "\n",
    "def C(B: float, D: bool) -> bool:\n",
    "    return B != D\n",
    "\n",
    "@tag(cache=\"pickle\")\n",
    "def E(C: bool, B: float) -> str:\n",
    "    return \"hello-world-ok\" * int(B)\n",
    "\n",
    "@tag(cache=\"json\")\n",
    "def F(E: str) -> dict:\n",
    "    return {E: E*3}\n",
    "\n",
    "@tag(cache=\"parquet\")\n",
    "def G() -> pd.DataFrame:\n",
    "    return pd.DataFrame({\"a\": [323, 3235], \"b\": [\"hello\", \"vorld\"]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hamilton import driver\n",
    "\n",
    "dr = (\n",
    "    driver.Builder()\n",
    "    .with_modules(primitives)\n",
    "    .with_cache()\n",
    "    .build()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['A', 'B', 'C', 'D', 'E', 'F', 'G'])\n"
     ]
    }
   ],
   "source": [
    "results = dr.execute(\n",
    "    [\"A\", \"B\", \"C\", \"D\", \"E\", \"F\", \"G\"],\n",
    "    inputs=dict(D=True),\n",
    "    overrides=dict(B=4)\n",
    ")\n",
    "print(results.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'A': 'eF4NxUEKwCAMBMCv9AlG48ZA6V+imx7t/28VBkaKJ5QqYa5LiXSzaMIMhuKULO2NAFnnmGDAZHQX9Hpc9/52Pj/UjxQx',\n",
       " 'C': 'eF4FwcERgDAIBMBWLCGAEDLj2MvpwTP2/3O3SztyWAgrUb46ZoNSTq2Zr5xkBWD+eEKMiykEsEJHhh7X/nbdP+gBFJk=',\n",
       " 'E': 'eF4FwckNgDAMBMBWKCHWOs5aQvQSfDxD/z9mAhT31VNFzDnENmoxM9L1taFMNfRcIk0HirsDbSQ6lPu6z3fq+QGUKRM7',\n",
       " 'F': 'eF4FwcERwCAIBMBWUgLincBMJr2Iwafp/5ddiIc3WanjTUQ3zUYO8wEIW9AU6LL7jF1SKZ5qoPoszlJe9/lOPT9v4hJQ',\n",
       " 'G': 'eF4FwckNgDAQA8BWKIEE7xEJ0YuJs8/Q/4+ZnvQQsoQl93yjsQuvZRhmBSVejiGiiXNwtThnucEJtzru/e31/NSBFG8='}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dr.cache.context_keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'metadata_store': ['initialized'],\n",
       " 'A': ['caching_behavior:resolved',\n",
       "  'code::versioned',\n",
       "  'in_memory_metadata::get::miss',\n",
       "  'context_key::created',\n",
       "  'metadata_store::get::hit',\n",
       "  'result_store::get::hit',\n",
       "  'context_key::created',\n",
       "  'metadata_store::get::hit',\n",
       "  'in_memory_metadata::set'],\n",
       " 'B': ['caching_behavior:resolved',\n",
       "  'code::versioned',\n",
       "  'as input',\n",
       "  'data::versioned',\n",
       "  'in_memory_metadata::set'],\n",
       " 'C': ['caching_behavior:resolved',\n",
       "  'code::versioned',\n",
       "  'in_memory_metadata::get::miss',\n",
       "  'context_key::created',\n",
       "  'metadata_store::get::hit',\n",
       "  'result_store::get::hit',\n",
       "  'context_key::created',\n",
       "  'metadata_store::get::hit',\n",
       "  'in_memory_metadata::set'],\n",
       " 'E': ['caching_behavior:resolved',\n",
       "  'code::versioned',\n",
       "  'in_memory_metadata::get::miss',\n",
       "  'context_key::created',\n",
       "  'metadata_store::get::hit',\n",
       "  'result_store::get::hit',\n",
       "  'context_key::created',\n",
       "  'metadata_store::get::hit',\n",
       "  'in_memory_metadata::set'],\n",
       " 'F': ['caching_behavior:resolved',\n",
       "  'code::versioned',\n",
       "  'in_memory_metadata::get::miss',\n",
       "  'context_key::created',\n",
       "  'metadata_store::get::hit',\n",
       "  'result_store::get::hit',\n",
       "  'context_key::created',\n",
       "  'metadata_store::get::hit',\n",
       "  'in_memory_metadata::set'],\n",
       " 'G': ['caching_behavior:resolved',\n",
       "  'code::versioned',\n",
       "  'in_memory_metadata::get::miss',\n",
       "  'context_key::created',\n",
       "  'metadata_store::get::hit',\n",
       "  'result_store::get::hit',\n",
       "  'context_key::created',\n",
       "  'metadata_store::get::hit',\n",
       "  'in_memory_metadata::set'],\n",
       " 'D': ['caching_behavior:resolved',\n",
       "  'code::versioned',\n",
       "  'as input',\n",
       "  'data::versioned',\n",
       "  'in_memory_metadata::set']}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dr.cache.logs(dr.cache.run_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Caching actions\n",
    "1. compute `code_version`\n",
    "2. compute `data_version` for `value` using `hashing_function`\n",
    "3. compute `value` by executing node with using `dependencies value`\n",
    "4. create `context_key` using `code_version` and `dependencies data_version`\n",
    "5. get `data_version` using `memory[node_name]`\n",
    "6. set `data_version` using `memory[node_name]`\n",
    "7. get `data_version` using `metadata_store[context_key]`\n",
    "8. set `data_version` using `metadata_store[context_key]`\n",
    "9. delete `data_version` using `metadata_store[context_key]`\n",
    "10. get `value` from `result_store[data_version]`\n",
    "11. set `value` from `result_store[data_version]`\n",
    "12. store `value` using a materializer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Power user\n",
    "- ignore dependencies from the `context_key` function\n",
    "- skip `set`/`get` operations for `metadata_store` and `result_store` \n",
    "- change the `hashing_function`\n",
    "- change the `metadata_store`\n",
    "- change the `result_store`\n",
    "- change the `context_key` function"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
